# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: leon

import os
from time import sleep
from absl import app, flags
import requests
import urllib.request
from xml.dom import minidom

FLAGS = flags.FLAGS


def download_csv(csv_url_prefix, csv_url, save_root):
  url = csv_url_prefix + csv_url
  file_name = csv_url.split('/')[-1]
  path = os.path.join(save_root, file_name)
  urllib.request.urlretrieve(url, path)


def parse_csv_urls(query_result):
  modified = query_result.split('\n', 2)[1]
  tree = minidom.parseString(modified)
  contents = tree.getElementsByTagName('Key')
  csv_urls = [content.firstChild.data for content in contents]
  return csv_urls


def query_csv_list_url(url):
  content = requests.get(url).content.decode()
  return content


def main(argv):
  csv_list_url = FLAGS.csv_list_url
  csv_url_prefix = FLAGS.csv_url_prefix
  save_root = FLAGS.save_root
  assert csv_list_url, 'csv_list_url must be specified.'
  assert csv_url_prefix, 'csv_url_prefix must be specified.'
  assert save_root, 'save_root must be specified.'

  query_result = query_csv_list_url(csv_list_url)
  csv_urls = parse_csv_urls(query_result)

  print('Total %d files.' % len(csv_urls))
  for csv_url in csv_urls:
    print('Downloading %s' % csv_url)
    download_csv(csv_url_prefix, csv_url, save_root)
    sleep(2)


if __name__ == '__main__':
  flags.DEFINE_string(
      'csv_list_url',
      'https://s3-eu-west-1.amazonaws.com/public.bitmex.com?delimiter=/&'
      'prefix=data/trade/&marker=data%2Ftrade%2F20170816.csv.gz',
      'Source csv file list url.')

  flags.DEFINE_string('csv_url_prefix',
                      'https://s3-eu-west-1.amazonaws.com/public.bitmex.com/',
                      'Source csv url prefix.')

  flags.DEFINE_string('save_root', None, 'Root dir to save download csvs.')

  app.run(main)
